{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dc5fea6a",
   "metadata": {},
   "source": [
    "# VTA 指令"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c122fed1",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "952f2444",
   "metadata": {},
   "outputs": [
    {
     "ename": "OSError",
     "evalue": "/media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/vta_hw/libvta_fsim_hw_vta2.0.so: undefined symbol: _ZN3vta4vmem20VirtualMemoryManager5AllocEm",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mOSError\u001b[39m                                   Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[2]\u001b[39m\u001b[32m, line 16\u001b[39m\n\u001b[32m      8\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mctypes\u001b[39;00m\n\u001b[32m     10\u001b[39m \u001b[38;5;66;03m# 先加载依赖库，设置RTLD_GLOBAL使符号全局可见\u001b[39;00m\n\u001b[32m     11\u001b[39m \u001b[38;5;66;03m# lib1 = ctypes.CDLL('/media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/vta_hw/libvta_fsim_hw_sim.so', ctypes.RTLD_GLOBAL)\u001b[39;00m\n\u001b[32m     12\u001b[39m \n\u001b[32m     13\u001b[39m \u001b[38;5;66;03m# 再加载第二个库，它可以访问第一个库中的符号\u001b[39;00m\n\u001b[32m     14\u001b[39m \u001b[38;5;66;03m# lib2 = ctypes.CDLL('/media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/libvta_fsim_sim.so')\u001b[39;00m\n\u001b[32m     15\u001b[39m \u001b[38;5;66;03m# lib = ctypes.CDLL(\"/media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/vta_hw/libvta_fsim_hw.so\", ctypes.RTLD_GLOBAL)\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m16\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mvta\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mtesting\u001b[39;00m\n\u001b[32m     17\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mvta\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mtesting\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m simulator\n\u001b[32m     19\u001b[39m np.random.seed(\u001b[32m0xDEADB\u001b[39m)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm-test/vta/python/vta/testing/__init__.py:20\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;66;03m# Licensed to the Apache Software Foundation (ASF) under one\u001b[39;00m\n\u001b[32m      2\u001b[39m \u001b[38;5;66;03m# or more contributor license agreements.  See the NOTICE file\u001b[39;00m\n\u001b[32m      3\u001b[39m \u001b[38;5;66;03m# distributed with this work for additional information\u001b[39;00m\n\u001b[32m   (...)\u001b[39m\u001b[32m     15\u001b[39m \u001b[38;5;66;03m# specific language governing permissions and limitations\u001b[39;00m\n\u001b[32m     16\u001b[39m \u001b[38;5;66;03m# under the License.\u001b[39;00m\n\u001b[32m     18\u001b[39m \u001b[33;03m\"\"\"Testing utilities, this namespace is not imported by default.\"\"\"\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m20\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01m.\u001b[39;00m\u001b[34;01mutils\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m run\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm-test/vta/python/vta/testing/utils.py:23\u001b[39m\n\u001b[32m     21\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtvm\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m rpc, autotvm\n\u001b[32m     22\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01m.\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01menvironment\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m get_env\n\u001b[32m---> \u001b[39m\u001b[32m23\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01m.\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m simulator\n\u001b[32m     26\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mrun\u001b[39m(run_func):\n\u001b[32m     27\u001b[39m \u001b[38;5;250m    \u001b[39m\u001b[33;03m\"\"\"Run test function on all available env.\u001b[39;00m\n\u001b[32m     28\u001b[39m \n\u001b[32m     29\u001b[39m \u001b[33;03m    Parameters\u001b[39;00m\n\u001b[32m     30\u001b[39m \u001b[33;03m    ----------\u001b[39;00m\n\u001b[32m     31\u001b[39m \u001b[33;03m    run_func : function(env, remote)\u001b[39;00m\n\u001b[32m     32\u001b[39m \u001b[33;03m    \"\"\"\u001b[39;00m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm-test/vta/python/vta/testing/simulator.py:119\u001b[39m\n\u001b[32m    110\u001b[39m \u001b[38;5;250m    \u001b[39m\u001b[33;03m\"\"\"Set debug mode\u001b[39;00m\n\u001b[32m    111\u001b[39m \u001b[33;03m    Paramaters\u001b[39;00m\n\u001b[32m    112\u001b[39m \u001b[33;03m    ----------\u001b[39;00m\n\u001b[32m    113\u001b[39m \u001b[33;03m    flag : int\u001b[39;00m\n\u001b[32m    114\u001b[39m \u001b[33;03m        The debug flag, 0 means clear all flags.\u001b[39;00m\n\u001b[32m    115\u001b[39m \u001b[33;03m    \"\"\"\u001b[39;00m\n\u001b[32m    116\u001b[39m     tvm.get_global_func(\u001b[33m\"\u001b[39m\u001b[33mvta.simulator.profiler_debug_mode\u001b[39m\u001b[33m\"\u001b[39m)(flag)\n\u001b[32m--> \u001b[39m\u001b[32m119\u001b[39m LIBS = \u001b[43m_load_sw\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm-test/vta/python/vta/testing/simulator.py:57\u001b[39m, in \u001b[36m_load_sw\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m     55\u001b[39m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[32m     56\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m require_sim:\n\u001b[32m---> \u001b[39m\u001b[32m57\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m err\n\u001b[32m     58\u001b[39m     warnings.warn(\u001b[33m\"\u001b[39m\u001b[33mError when loading VTA driver \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[33m: \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[33m\"\u001b[39m.format(lib_driver[\u001b[32m0\u001b[39m], err))\n\u001b[32m     59\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m []\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/board/arria10/lxw/tasks/tvm-test/vta/python/vta/testing/simulator.py:51\u001b[39m, in \u001b[36m_load_sw\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m     48\u001b[39m lib_runtime = find_libvta(lib_driver_name.replace(\u001b[33m\"\u001b[39m\u001b[33m_hw\u001b[39m\u001b[33m\"\u001b[39m, \u001b[33m\"\u001b[39m\u001b[33m\"\u001b[39m), optional=(\u001b[38;5;129;01mnot\u001b[39;00m require_sim))\n\u001b[32m     50\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m---> \u001b[39m\u001b[32m51\u001b[39m     lib_driver = \u001b[43mctypes\u001b[49m\u001b[43m.\u001b[49m\u001b[43mCDLL\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlib_driver\u001b[49m\u001b[43m[\u001b[49m\u001b[32;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mctypes\u001b[49m\u001b[43m.\u001b[49m\u001b[43mRTLD_GLOBAL\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     52\u001b[39m     lib_runtime = ctypes.CDLL(lib_runtime[\u001b[32m0\u001b[39m], ctypes.RTLD_GLOBAL)\n\u001b[32m     54\u001b[39m     libs = [lib_driver, lib_runtime]\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/media/pc/data/lxw/envs/anaconda3a/envs/xin/lib/python3.12/ctypes/__init__.py:379\u001b[39m, in \u001b[36mCDLL.__init__\u001b[39m\u001b[34m(self, name, mode, handle, use_errno, use_last_error, winmode)\u001b[39m\n\u001b[32m    376\u001b[39m \u001b[38;5;28mself\u001b[39m._FuncPtr = _FuncPtr\n\u001b[32m    378\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m handle \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m379\u001b[39m     \u001b[38;5;28mself\u001b[39m._handle = \u001b[43m_dlopen\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    380\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    381\u001b[39m     \u001b[38;5;28mself\u001b[39m._handle = handle\n",
      "\u001b[31mOSError\u001b[39m: /media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/vta_hw/libvta_fsim_hw_vta2.0.so: undefined symbol: _ZN3vta4vmem20VirtualMemoryManager5AllocEm"
     ]
    }
   ],
   "source": [
    "import tvm\n",
    "from tvm import te\n",
    "import numpy as np\n",
    "from tvm import topi\n",
    "from tvm.contrib import utils\n",
    "\n",
    "import vta\n",
    "import ctypes\n",
    "\n",
    "# 先加载依赖库，设置RTLD_GLOBAL使符号全局可见\n",
    "# lib1 = ctypes.CDLL('/media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/vta_hw/libvta_fsim_hw_sim.so', ctypes.RTLD_GLOBAL)\n",
    "\n",
    "# 再加载第二个库，它可以访问第一个库中的符号\n",
    "# lib2 = ctypes.CDLL('/media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/libvta_fsim_sim.so')\n",
    "# lib = ctypes.CDLL(\"/media/pc/data/board/arria10/lxw/tasks/tvm-test/build/vta/vta_hw/libvta_fsim_hw.so\", ctypes.RTLD_GLOBAL)\n",
    "import vta.testing\n",
    "from vta.testing import simulator\n",
    "\n",
    "np.random.seed(0xDEADB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18762a92",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_runtime_array():\n",
    "    def _run(env, remote):\n",
    "        n = 100\n",
    "        dev = tvm.ext_dev(0)\n",
    "        x_np = np.random.randint(1, 10, size=(n, n, env.BATCH, env.BLOCK_OUT)).astype(\"int8\")\n",
    "        x_nd = tvm.nd.array(x_np, dev)\n",
    "        np.testing.assert_equal(x_np, x_nd.numpy())\n",
    "\n",
    "    vta.testing.run(_run)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4040f9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_save_load_out():\n",
    "    \"\"\"Test save/store output command\"\"\"\n",
    "\n",
    "    def _run(env, remote):\n",
    "        n = 6\n",
    "        x = te.placeholder((n, n, env.BATCH, env.BLOCK_OUT), name=\"x\", dtype=env.acc_dtype)\n",
    "        x_buf = te.compute((n, n, env.BATCH, env.BLOCK_OUT), lambda *i: x(*i), \"x_buf\")\n",
    "        # insert no-op that won't be optimized away\n",
    "        y_buf = te.compute((n, n, env.BATCH, env.BLOCK_OUT), lambda *i: x_buf(*i) >> 0, \"y_buf\")\n",
    "        y = te.compute(\n",
    "            (n, n, env.BATCH, env.BLOCK_OUT), lambda *i: y_buf(*i).astype(env.inp_dtype), \"y\"\n",
    "        )\n",
    "        # schedule\n",
    "        s = te.create_schedule(y.op)\n",
    "        s[x_buf].set_scope(env.acc_scope)\n",
    "        s[x_buf].pragma(x_buf.op.axis[0], env.dma_copy)\n",
    "        s[y_buf].set_scope(env.acc_scope)\n",
    "        s[y_buf].pragma(y_buf.op.axis[0], env.alu)\n",
    "        s[y].pragma(y.op.axis[0], env.dma_copy)\n",
    "\n",
    "        # verification\n",
    "        with vta.build_config():\n",
    "            m = vta.build(s, [x, y], tvm.target.Target(\"ext_dev\", host=env.target_host))\n",
    "\n",
    "        if not remote:\n",
    "            return\n",
    "        temp = utils.tempdir()\n",
    "        m.save(temp.relpath(\"load_act.o\"))\n",
    "        remote.upload(temp.relpath(\"load_act.o\"))\n",
    "        f = remote.load_module(\"load_act.o\")\n",
    "        # verify\n",
    "        dev = tvm.ext_dev(0)\n",
    "        x_np = np.random.randint(1, 10, size=(n, n, env.BATCH, env.BLOCK_OUT)).astype(x.dtype)\n",
    "        y_np = x_np.astype(y.dtype)\n",
    "        x_nd = tvm.nd.array(x_np, dev)\n",
    "        y_nd = tvm.nd.empty(y_np.shape, device=dev, dtype=y_np.dtype)\n",
    "\n",
    "        if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "            simulator.clear_stats()\n",
    "\n",
    "        f(x_nd, y_nd)\n",
    "\n",
    "        np.testing.assert_equal(y_np, y_nd.numpy())\n",
    "\n",
    "        if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "            sim_stats = simulator.stats()\n",
    "            print(\"Save load execution statistics:\")\n",
    "            for k, v in sim_stats.items():\n",
    "                print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "def test_padded_load():\n",
    "    \"\"\"Test padded load.\"\"\"\n",
    "\n",
    "    def _run(env, remote):\n",
    "        def check_padded_load(pad_before, pad_after, test_name=None):\n",
    "            # declare\n",
    "            n = 3\n",
    "            m = 5\n",
    "            x = te.placeholder((n, m, env.BATCH, env.BLOCK_OUT), name=\"x\", dtype=env.acc_dtype)\n",
    "            x_buf = topi.nn.pad(x, pad_before, pad_after, name=\"y\")\n",
    "            # insert no-op that won't be optimized away\n",
    "            y_buf = te.compute(\n",
    "                (\n",
    "                    n + pad_before[0] + pad_after[0],\n",
    "                    m + pad_before[1] + pad_after[1],\n",
    "                    env.BATCH,\n",
    "                    env.BLOCK_OUT,\n",
    "                ),\n",
    "                lambda *i: x_buf(*i) >> 0,\n",
    "                \"y_buf\",\n",
    "            )\n",
    "            y = te.compute(\n",
    "                (\n",
    "                    n + pad_before[0] + pad_after[0],\n",
    "                    m + pad_before[1] + pad_after[1],\n",
    "                    env.BATCH,\n",
    "                    env.BLOCK_OUT,\n",
    "                ),\n",
    "                lambda *i: y_buf(*i).astype(env.inp_dtype),\n",
    "                \"y\",\n",
    "            )\n",
    "            # schedule\n",
    "            s = te.create_schedule(y.op)\n",
    "            s[x_buf].set_scope(env.acc_scope)\n",
    "            s[x_buf].pragma(x_buf.op.axis[0], env.dma_copy)\n",
    "            s[y_buf].set_scope(env.acc_scope)\n",
    "            s[y_buf].pragma(y_buf.op.axis[0], env.alu)\n",
    "            s[y].pragma(y.op.axis[0], env.dma_copy)\n",
    "            # build\n",
    "            with vta.build_config():\n",
    "                mod = vta.build(s, [x, y], tvm.target.Target(\"ext_dev\", host=env.target_host))\n",
    "\n",
    "            if not remote:\n",
    "                return\n",
    "            temp = utils.tempdir()\n",
    "            mod.save(temp.relpath(\"padded_load.o\"))\n",
    "            remote.upload(temp.relpath(\"padded_load.o\"))\n",
    "            f = remote.load_module(\"padded_load.o\")\n",
    "            # verify\n",
    "            dev = remote.ext_dev(0)\n",
    "            x_np = np.random.randint(0, 10, size=(n, m, env.BATCH, env.BLOCK_OUT)).astype(x.dtype)\n",
    "            y_np = np.zeros(\n",
    "                (\n",
    "                    n + pad_before[0] + pad_after[0],\n",
    "                    m + pad_before[1] + pad_after[1],\n",
    "                    env.BATCH,\n",
    "                    env.BLOCK_OUT,\n",
    "                )\n",
    "            ).astype(y.dtype)\n",
    "            y_np[pad_before[0] : pad_before[0] + n, pad_before[1] : pad_before[1] + m, :] = x_np\n",
    "            x_nd = tvm.nd.array(x_np, dev)\n",
    "            y_nd = tvm.nd.empty(y_np.shape, device=dev, dtype=y_np.dtype)\n",
    "\n",
    "            if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "                simulator.clear_stats()\n",
    "\n",
    "            f(x_nd, y_nd)\n",
    "\n",
    "            np.testing.assert_equal(y_np, y_nd.numpy())\n",
    "\n",
    "            if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "                sim_stats = simulator.stats()\n",
    "                print(\"Padded {} load execution statistics:\".format(test_name))\n",
    "                for k, v in sim_stats.items():\n",
    "                    print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
    "\n",
    "        check_padded_load([2, 0, 0, 0], [0, 0, 0, 0], test_name=\"Y0\")\n",
    "        check_padded_load([0, 2, 0, 0], [0, 0, 0, 0], test_name=\"Y1\")\n",
    "        check_padded_load([0, 0, 0, 0], [2, 0, 0, 0], test_name=\"X0\")\n",
    "        check_padded_load([0, 0, 0, 0], [0, 2, 0, 0], test_name=\"X1\")\n",
    "        check_padded_load([1, 1, 0, 0], [1, 1, 0, 0], test_name=\"all\")\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "def test_gemm():\n",
    "    \"\"\"Test GEMM.\"\"\"\n",
    "\n",
    "    def _run(env, remote):\n",
    "        # declare\n",
    "        o = 4\n",
    "        n = 1\n",
    "        m = 4\n",
    "        x = te.placeholder((o, n, env.BATCH, env.BLOCK_IN), name=\"x\", dtype=env.inp_dtype)\n",
    "        w = te.placeholder((m, n, env.BLOCK_OUT, env.BLOCK_IN), name=\"w\", dtype=env.wgt_dtype)\n",
    "        x_buf = te.compute((o, n, env.BATCH, env.BLOCK_IN), lambda *i: x(*i), \"x_buf\")\n",
    "        w_buf = te.compute((m, n, env.BLOCK_OUT, env.BLOCK_IN), lambda *i: w(*i), \"w_buf\")\n",
    "        ko = te.reduce_axis((0, n), name=\"ko\")\n",
    "        ki = te.reduce_axis((0, env.BLOCK_IN), name=\"ki\")\n",
    "        y_gem = te.compute(\n",
    "            (o, m, env.BATCH, env.BLOCK_OUT),\n",
    "            lambda bo, co, bi, ci: te.sum(\n",
    "                x_buf[bo, ko, bi, ki].astype(env.acc_dtype)\n",
    "                * w_buf[co, ko, ci, ki].astype(env.acc_dtype),\n",
    "                axis=[ko, ki],\n",
    "            ),\n",
    "            name=\"y_gem\",\n",
    "        )\n",
    "        y_shf = te.compute(\n",
    "            (o, m, env.BATCH, env.BLOCK_OUT), lambda *i: y_gem(*i) >> 8, name=\"y_shf\"\n",
    "        )\n",
    "        y_max = te.compute(\n",
    "            (o, m, env.BATCH, env.BLOCK_OUT), lambda *i: tvm.te.max(y_shf(*i), 0), \"y_max\"\n",
    "        )  # relu\n",
    "        y_min = te.compute(\n",
    "            (o, m, env.BATCH, env.BLOCK_OUT),\n",
    "            lambda *i: tvm.te.min(y_max(*i), (1 << (env.INP_WIDTH - 1)) - 1),\n",
    "            \"y_min\",\n",
    "        )  # relu\n",
    "        y = te.compute(\n",
    "            (o, m, env.BATCH, env.BLOCK_OUT), lambda *i: y_min(*i).astype(env.inp_dtype), name=\"y\"\n",
    "        )\n",
    "\n",
    "        if not remote:\n",
    "            return\n",
    "\n",
    "        def verify(s, name=None):\n",
    "            # Build with the CSE pass disabled as otherwise it would complicate the test\n",
    "            with vta.build_config(disabled_pass={\"tir.CommonSubexprElimTIR\"}):\n",
    "                mod = vta.build(s, [x, w, y], tvm.target.Target(\"ext_dev\", host=env.target_host))\n",
    "            temp = utils.tempdir()\n",
    "            mod.save(temp.relpath(\"gemm.o\"))\n",
    "            remote.upload(temp.relpath(\"gemm.o\"))\n",
    "            f = remote.load_module(\"gemm.o\")\n",
    "            # verify\n",
    "            dev = remote.ext_dev(0)\n",
    "            x_np = np.random.randint(-128, 128, size=(o, n, env.BATCH, env.BLOCK_IN)).astype(\n",
    "                x.dtype\n",
    "            )\n",
    "            w_np = np.random.randint(-128, 128, size=(m, n, env.BLOCK_OUT, env.BLOCK_IN)).astype(\n",
    "                w.dtype\n",
    "            )\n",
    "            y_np = np.zeros((o, m, env.BATCH, env.BLOCK_OUT)).astype(y.dtype)\n",
    "            x_nd = tvm.nd.array(x_np, dev)\n",
    "            w_nd = tvm.nd.array(w_np, dev)\n",
    "            y_nd = tvm.nd.array(y_np, dev)\n",
    "            y_np = y_np.astype(env.acc_dtype)\n",
    "            for b in range(o):\n",
    "                for i in range(m):\n",
    "                    for j in range(n):\n",
    "                        y_np[b, i, :] += np.dot(\n",
    "                            x_np[b, j, :].astype(env.acc_dtype), w_np[i, j].T.astype(env.acc_dtype)\n",
    "                        )\n",
    "            y_np = np.right_shift(y_np, 8)\n",
    "            y_np = np.clip(y_np, 0, (1 << (env.INP_WIDTH - 1)) - 1).astype(y.dtype)\n",
    "\n",
    "            if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "                simulator.clear_stats()\n",
    "\n",
    "            f(x_nd, w_nd, y_nd)\n",
    "\n",
    "            np.testing.assert_equal(y_np, y_nd.numpy())\n",
    "\n",
    "            if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "                sim_stats = simulator.stats()\n",
    "                print(\"GEMM schedule:{} execution statistics:\".format(name))\n",
    "                for k, v in sim_stats.items():\n",
    "                    print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
    "\n",
    "        def test_schedule1():\n",
    "            # default schedule with no smt\n",
    "            s = te.create_schedule(y.op)\n",
    "            # set the scope of the SRAM buffers\n",
    "            s[x_buf].set_scope(env.inp_scope)\n",
    "            s[w_buf].set_scope(env.wgt_scope)\n",
    "            s[y_gem].set_scope(env.acc_scope)\n",
    "            s[y_shf].set_scope(env.acc_scope)\n",
    "            s[y_max].set_scope(env.acc_scope)\n",
    "            s[y_min].set_scope(env.acc_scope)\n",
    "            # set pragmas for DMA transfer and ALU ops\n",
    "            s[x_buf].compute_at(s[y_gem], ko)\n",
    "            s[x_buf].pragma(s[x_buf].op.axis[0], env.dma_copy)\n",
    "            s[w_buf].compute_at(s[y_gem], ko)\n",
    "            s[w_buf].pragma(s[w_buf].op.axis[0], env.dma_copy)\n",
    "            s[y_shf].pragma(s[y_shf].op.axis[0], env.alu)\n",
    "            s[y_max].pragma(s[y_max].op.axis[0], env.alu)\n",
    "            s[y_min].pragma(s[y_min].op.axis[0], env.alu)\n",
    "            s[y].pragma(s[y].op.axis[0], env.dma_copy)\n",
    "            # tensorization\n",
    "            s[y_gem].reorder(\n",
    "                ko,\n",
    "                s[y_gem].op.axis[0],\n",
    "                s[y_gem].op.axis[1],\n",
    "                s[y_gem].op.axis[2],\n",
    "                s[y_gem].op.axis[3],\n",
    "                ki,\n",
    "            )\n",
    "            s[y_gem].tensorize(s[y_gem].op.axis[2], env.gemm)\n",
    "            verify(s, name=\"default\")\n",
    "\n",
    "        def test_smt():\n",
    "            # test smt schedule\n",
    "            s = te.create_schedule(y.op)\n",
    "            s[x_buf].set_scope(env.inp_scope)\n",
    "            s[w_buf].set_scope(env.wgt_scope)\n",
    "            s[y_gem].set_scope(env.acc_scope)\n",
    "            s[y_shf].set_scope(env.acc_scope)\n",
    "            s[y_max].set_scope(env.acc_scope)\n",
    "            s[y_min].set_scope(env.acc_scope)\n",
    "            abo, aco, abi, aci = s[y].op.axis\n",
    "            abo1, abo2 = s[y].split(abo, nparts=2)\n",
    "            s[y].bind(abo1, te.thread_axis(\"cthread\"))\n",
    "            s[y_gem].compute_at(s[y], abo1)\n",
    "            s[y_shf].compute_at(s[y], abo1)\n",
    "            s[y_max].compute_at(s[y], abo1)\n",
    "            s[y_min].compute_at(s[y], abo1)\n",
    "            s[y_gem].reorder(\n",
    "                ko,\n",
    "                s[y_gem].op.axis[0],\n",
    "                s[y_gem].op.axis[1],\n",
    "                s[y_gem].op.axis[2],\n",
    "                s[y_gem].op.axis[3],\n",
    "                ki,\n",
    "            )\n",
    "            s[y_gem].tensorize(s[y_gem].op.axis[2], env.gemm)\n",
    "            s[y_shf].pragma(s[y_shf].op.axis[0], env.alu)\n",
    "            s[y_max].pragma(s[y_max].op.axis[0], env.alu)\n",
    "            s[y_min].pragma(s[y_min].op.axis[0], env.alu)\n",
    "            s[x_buf].compute_at(s[y_gem], ko)\n",
    "            s[x_buf].pragma(s[x_buf].op.axis[0], env.dma_copy)\n",
    "            s[w_buf].compute_at(s[y_gem], ko)\n",
    "            s[w_buf].pragma(s[w_buf].op.axis[0], env.dma_copy)\n",
    "            s[y].pragma(abo2, env.dma_copy)\n",
    "            verify(s, name=\"smt\")\n",
    "\n",
    "        test_schedule1()\n",
    "        test_smt()\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "def test_alu():\n",
    "    def _run(env, remote):\n",
    "        def check_alu(tvm_op, np_op=None, use_imm=False, test_name=None):\n",
    "            \"\"\"Test ALU\"\"\"\n",
    "            m = 8\n",
    "            n = 8\n",
    "            imm = np.random.randint(1, 5)\n",
    "            # compute\n",
    "            a = te.placeholder((m, n, env.BATCH, env.BLOCK_OUT), name=\"a\", dtype=env.acc_dtype)\n",
    "            a_buf = te.compute(\n",
    "                (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: a(*i), \"a_buf\"\n",
    "            )  # DRAM->SRAM\n",
    "            if use_imm:\n",
    "                res_buf = te.compute(\n",
    "                    (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: tvm_op(a_buf(*i), imm), \"res_buf\"\n",
    "                )  # compute\n",
    "            else:\n",
    "                b = te.placeholder((m, n, env.BATCH, env.BLOCK_OUT), name=\"b\", dtype=env.acc_dtype)\n",
    "                b_buf = te.compute(\n",
    "                    (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: b(*i), \"b_buf\"\n",
    "                )  # DRAM->SRAM\n",
    "                res_buf = te.compute(\n",
    "                    (m, n, env.BATCH, env.BLOCK_OUT),\n",
    "                    lambda *i: tvm_op(a_buf(*i), b_buf(*i)),\n",
    "                    \"res_buf\",\n",
    "                )  # compute5B\n",
    "            res = te.compute(\n",
    "                (m, n, env.BATCH, env.BLOCK_OUT),\n",
    "                lambda *i: res_buf(*i).astype(env.inp_dtype),\n",
    "                \"res\",\n",
    "            )  # SRAM->DRAM\n",
    "            # schedule\n",
    "            s = te.create_schedule(res.op)\n",
    "            s[a_buf].set_scope(env.acc_scope)  # SRAM\n",
    "            s[a_buf].pragma(a_buf.op.axis[0], env.dma_copy)  # DRAM->SRAM\n",
    "            s[res_buf].set_scope(env.acc_scope)  # SRAM\n",
    "            s[res_buf].pragma(res_buf.op.axis[0], env.alu)  # compute\n",
    "            s[res].pragma(res.op.axis[0], env.dma_copy)  # SRAM->DRAM\n",
    "            if not use_imm:\n",
    "                s[b_buf].set_scope(env.acc_scope)  # SRAM\n",
    "                s[b_buf].pragma(b_buf.op.axis[0], env.dma_copy)  # DRAM->SRAM\n",
    "\n",
    "            if not remote:\n",
    "                return\n",
    "\n",
    "            # build\n",
    "            with vta.build_config():\n",
    "                if use_imm:\n",
    "                    mod = vta.build(s, [a, res], tvm.target.Target(\"ext_dev\", host=env.target_host))\n",
    "                else:\n",
    "                    mod = vta.build(\n",
    "                        s, [a, b, res], tvm.target.Target(\"ext_dev\", host=env.target_host)\n",
    "                    )\n",
    "            temp = utils.tempdir()\n",
    "            mod.save(temp.relpath(\"load_act.o\"))\n",
    "            remote.upload(temp.relpath(\"load_act.o\"))\n",
    "            f = remote.load_module(\"load_act.o\")\n",
    "            # verify\n",
    "            dev = remote.ext_dev(0)\n",
    "            a_np = np.random.randint(-16, 16, size=(m, n, env.BATCH, env.BLOCK_OUT)).astype(a.dtype)\n",
    "            if use_imm:\n",
    "                res_np = np_op(a_np, imm) if np_op else tvm_op(a_np, imm)\n",
    "            else:\n",
    "                b_np = np.random.randint(-16, 16, size=(m, n, env.BATCH, env.BLOCK_OUT)).astype(\n",
    "                    b.dtype\n",
    "                )\n",
    "                res_np = np_op(a_np, b_np) if np_op else tvm_op(a_np, b_np)\n",
    "            res_np = res_np.astype(res.dtype)\n",
    "            a_nd = tvm.nd.array(a_np, dev)\n",
    "            res_nd = tvm.nd.array(np.zeros((m, n, env.BATCH, env.BLOCK_OUT)).astype(res.dtype), dev)\n",
    "\n",
    "            if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "                simulator.clear_stats()\n",
    "\n",
    "            if use_imm:\n",
    "                f(a_nd, res_nd)\n",
    "            else:\n",
    "                b_nd = tvm.nd.array(b_np, dev)\n",
    "                f(a_nd, b_nd, res_nd)\n",
    "\n",
    "            np.testing.assert_equal(res_np, res_nd.numpy())\n",
    "\n",
    "            if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "                sim_stats = simulator.stats()\n",
    "                print(\"ALU {} execution statistics:\".format(test_name))\n",
    "                for k, v in sim_stats.items():\n",
    "                    print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
    "\n",
    "        check_alu(lambda x, y: x << y, np.left_shift, use_imm=True, test_name=\"SHL\")\n",
    "        check_alu(tvm.te.max, np.maximum, use_imm=True, test_name=\"MAX\")\n",
    "        check_alu(tvm.te.max, np.maximum, test_name=\"MAX\")\n",
    "        check_alu(lambda x, y: x + y, use_imm=True, test_name=\"ADD\")\n",
    "        check_alu(lambda x, y: x + y, test_name=\"ADD\")\n",
    "        check_alu(lambda x, y: x >> y, np.right_shift, use_imm=True, test_name=\"SHR\")\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "def test_relu():\n",
    "    \"\"\"Test RELU on ALU\"\"\"\n",
    "\n",
    "    def _run(env, remote):\n",
    "        m = 8\n",
    "        n = 10\n",
    "        # compute\n",
    "        a = te.placeholder((m, n, env.BATCH, env.BLOCK_OUT), name=\"a\", dtype=env.acc_dtype)\n",
    "        a_buf = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: a(*i), \"a_buf\"\n",
    "        )  # DRAM->SRAM\n",
    "        max_buf = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: tvm.te.max(a_buf(*i), 0), \"res_buf\"\n",
    "        )  # relu\n",
    "        min_buf = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT),\n",
    "            lambda *i: tvm.te.min(max_buf(*i), (1 << (env.INP_WIDTH - 1)) - 1),\n",
    "            \"max_buf\",\n",
    "        )  # relu\n",
    "        res = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT),\n",
    "            lambda *i: min_buf(*i).astype(env.inp_dtype),\n",
    "            \"min_buf\",\n",
    "        )  # SRAM->DRAM\n",
    "        # schedule\n",
    "        s = te.create_schedule(res.op)\n",
    "        s[a_buf].set_scope(env.acc_scope)  # SRAM\n",
    "        s[a_buf].pragma(a_buf.op.axis[0], env.dma_copy)  # DRAM->SRAM\n",
    "        s[max_buf].set_scope(env.acc_scope)  # SRAM\n",
    "        s[min_buf].set_scope(env.acc_scope)  # SRAM\n",
    "        s[max_buf].pragma(max_buf.op.axis[0], env.alu)  # compute\n",
    "        s[min_buf].pragma(min_buf.op.axis[0], env.alu)  # compute\n",
    "        s[res].pragma(res.op.axis[0], env.dma_copy)  # SRAM->DRAM\n",
    "        # build\n",
    "        with vta.build_config():\n",
    "            mod = vta.build(s, [a, res], tvm.target.Target(\"ext_dev\", host=env.target_host))\n",
    "        if not remote:\n",
    "            return\n",
    "        temp = utils.tempdir()\n",
    "        mod.save(temp.relpath(\"load_act.o\"))\n",
    "        remote.upload(temp.relpath(\"load_act.o\"))\n",
    "        f = remote.load_module(\"load_act.o\")\n",
    "        # verify\n",
    "        dev = remote.ext_dev(0)\n",
    "        a_np = np.random.randint(-256, 256, size=(m, n, env.BATCH, env.BLOCK_OUT)).astype(a.dtype)\n",
    "        res_np = np.clip(a_np, 0, (1 << (env.INP_WIDTH - 1)) - 1).astype(res.dtype)\n",
    "        a_nd = tvm.nd.array(a_np, dev)\n",
    "        res_nd = tvm.nd.array(np.zeros((m, n, env.BATCH, env.BLOCK_OUT)).astype(res.dtype), dev)\n",
    "\n",
    "        if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "            simulator.clear_stats()\n",
    "\n",
    "        f(a_nd, res_nd)\n",
    "\n",
    "        np.testing.assert_equal(res_np, res_nd.numpy())\n",
    "\n",
    "        if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "            sim_stats = simulator.stats()\n",
    "            print(\"Relu execution statistics:\")\n",
    "            for k, v in sim_stats.items():\n",
    "                print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "def test_shift_and_scale():\n",
    "    \"\"\"Test shift and scale on ALU\"\"\"\n",
    "\n",
    "    def _run(env, remote):\n",
    "        m = 2\n",
    "        n = 8\n",
    "        imm_shift = np.random.randint(0, 8)\n",
    "        imm_scale = np.random.randint(1, 5)\n",
    "        # compute\n",
    "        a = te.placeholder((m, n, env.BATCH, env.BLOCK_OUT), name=\"a\", dtype=env.acc_dtype)\n",
    "        a_buf = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: a(*i), \"a_buf\"\n",
    "        )  # DRAM->SRAM\n",
    "        res_shift = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: a_buf(*i) + imm_shift, \"res_shift\"\n",
    "        )  # compute\n",
    "        res_scale = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: res_shift(*i) >> imm_scale, \"res_scale\"\n",
    "        )  # compute\n",
    "        res = te.compute(\n",
    "            (m, n, env.BATCH, env.BLOCK_OUT), lambda *i: res_scale(*i).astype(env.inp_dtype), \"res\"\n",
    "        )  # SRAM->DRAM\n",
    "        # schedule\n",
    "        s = te.create_schedule(res.op)\n",
    "        s[a_buf].set_scope(env.acc_scope)  # SRAM\n",
    "        s[res_shift].set_scope(env.acc_scope)  # SRAM\n",
    "        s[res_scale].set_scope(env.acc_scope)  # SRAM\n",
    "        s[a_buf].pragma(a_buf.op.axis[0], env.dma_copy)  # DRAM->SRAM\n",
    "        s[res_shift].pragma(res_shift.op.axis[0], env.alu)  # compute\n",
    "        s[res_scale].pragma(res_scale.op.axis[0], env.alu)  # compute\n",
    "        s[res].pragma(res.op.axis[0], env.dma_copy)  # SRAM->DRAM\n",
    "        # build\n",
    "        mod = vta.build(s, [a, res], tvm.target.Target(\"ext_dev\", host=env.target_host))\n",
    "        if not remote:\n",
    "            return\n",
    "        temp = utils.tempdir()\n",
    "        mod.save(temp.relpath(\"load_act.o\"))\n",
    "        remote.upload(temp.relpath(\"load_act.o\"))\n",
    "        f = remote.load_module(\"load_act.o\")\n",
    "        # verify\n",
    "        dev = remote.ext_dev(0)\n",
    "        a_np = np.random.randint(-10, 10, size=(m, n, env.BATCH, env.BLOCK_OUT)).astype(a.dtype)\n",
    "        res_np = np.right_shift((a_np + imm_shift), imm_scale)\n",
    "        res_np = res_np.astype(res.dtype)\n",
    "        a_nd = tvm.nd.array(a_np, dev)\n",
    "        res_nd = tvm.nd.array(np.zeros((m, n, env.BATCH, env.BLOCK_OUT)).astype(res.dtype), dev)\n",
    "\n",
    "        if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "            simulator.clear_stats()\n",
    "\n",
    "        f(a_nd, res_nd)\n",
    "\n",
    "        np.testing.assert_equal(res_np, res_nd.numpy())\n",
    "\n",
    "        if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "            sim_stats = simulator.stats()\n",
    "            print(\"Shift and scale execution statistics:\")\n",
    "            for k, v in sim_stats.items():\n",
    "                print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "def test_runtime_array():\n",
    "    def _run(env, remote):\n",
    "        n = 100\n",
    "        dev = remote.ext_dev(0)\n",
    "        x_np = np.random.randint(1, 10, size=(n, n, env.BATCH, env.BLOCK_OUT)).astype(\"int8\")\n",
    "        x_nd = tvm.nd.array(x_np, dev)\n",
    "        np.testing.assert_equal(x_np, x_nd.numpy())\n",
    "\n",
    "    vta.testing.run(_run)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c6f6d4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_runtime_array()\n",
    "test_save_load_out()\n",
    "test_padded_load()\n",
    "test_gemm()\n",
    "test_alu()\n",
    "test_relu()\n",
    "test_shift_and_scale()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e7cb194",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
